#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Time    : 2022/3/20 1:20 下午
# @Author  : WangZhixing
import os
import shutil
import sys
from Visualization.Visualize import visualize

rootPath = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
sys.path.append(rootPath)

import argparse
from torch_geometric.transforms import NormalizeFeatures
from ProcessData import DependenceGraph
from Utils import ConfigFile
from Model.Module.GraphSage import GraphSage
import pandas as pd
import torch
from Metric import Metric
from Output.output_mehod.result2rsf_file import result2rsf_file


def train(model, data, criterion, optimizer):
    model.train()
    optimizer.zero_grad()  # 清除梯度
    out = model(data.x, data.edge_index)  # Perform a single forward pass.执行单次向前传播
    loss = criterion(out[data.train_mask], data.y[data.train_mask])  # 根据训练节点计算损失
    # Compute the loss solely based on the training nodes.
    loss.backward()  # Derive gradients，获取梯度
    optimizer.step()  # Update parameters based on gradients.根据梯度更新参数
    return loss

def test(model, data):
    model.eval()
    out = model(data.x, data.edge_index)
    pred = out.argmax(dim=1)  # Use the class with highest probability.
    # 根据实际的标签进行检查
    test_correct = pred[data.test_mask] == data.y[data.test_mask]  # Check against ground-truth labels.
    # 得出预测正确的比例
    test_acc = int(test_correct.sum()) / int(data.test_mask.sum())  # Derive ratio of correct predictions.
    # 返回预测正确的比例
    return test_acc

def GraphSage_train(data,**kwarg):
    model = GraphSage(
        num_features=data.num_features,
        hidden_channels=kwarg['model_layer'][0],
        num_classes=kwarg['model_layer'][1]
    )

    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=kwarg['lr'], weight_decay=5e-4)  # 定义Adam优化器
    df = pd.DataFrame(columns=["Loss"])
    df.index.name = "Epoch"
    for epoch in range(1, kwarg['train_epoch']):
        loss = train(model, data, criterion, optimizer)
    test_acc = test(model, data)
    print(f'Test Accuracy: {test_acc:.4f}')
    model.eval()
    out = model(data.x, data.edge_index)
    preds = out.argmax(dim=1).tolist()
    return out, preds




if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="ReadConfig")
    parser.add_argument("-c", "--config", type=str)
    args = parser.parse_args()
    kwarg = ConfigFile(args.config).ReadConfig()
    if os.path.exists(os.path.join(kwarg["root"],"processed")):
        shutil.rmtree(os.path.join(kwarg["root"],"processed"))
    data = DependenceGraph(kwarg["root"], transform=NormalizeFeatures(),).data
    out, preds=GraphSage_train(data,**kwarg)
    result2rsf_file(kwarg["root"], preds, kwarg["outfile_path"])
    Metric(kwarg["project"], kwarg["outfile_path"], kwarg["ground_path"],dep_file=os.path.join(kwarg["root"], "raw","edge.rsf"))
    visualize(out,preds)
